import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap

# 记录种类及数量


def count(column):
    column_type = {}
    for obj in column:
        if obj in column_type.keys():
            column_type[obj] = column_type[obj] + 1
        else:
            column_type[obj] = 0
    temp = column_type.copy()
    column_type.clear()
    # 字典
    column_type = {k: v for k, v in sorted(
        temp.items(), key=lambda item: item[1], reverse=True)}
    return column_type

# 柱状图


def draw_bar(data, column, dir):
    column_type = count(data[column])
    labels = list(column_type.keys())
    sizes = list(column_type.values())
    # 曲线救国...
    for i in range(len(labels)):
        if labels[i] == '上海市':
            labels[i] = 'Shanghai'
        if labels[i] == '北京市':
            labels[i] = 'Beijing'
    plt.figure(figsize=(20, 15), dpi=80)
    # 防止种类过多只取前8个
    if len(column_type) > 8:
        labels = labels[:8]
        sizes = sizes[:8]
    plt.title('numbers of ' + column, fontsize=16)
    for a, b in zip(labels, sizes):
        plt.text(a, b + 0.05, '%.0f' %
                 b, ha='center', va='bottom', fontsize=8)
    plt.bar(labels, sizes, color="#4CAF50", width=0.5)
    plt.savefig(r'./' + dir + r'/' +
                column.replace('/', 'Or') + '.png')

# 饼状图


def draw_pie(data, column_name, dir):
    column_type = count(data[column_name])
    if len(column_type) > 7:
        temp = column_type.copy()
        column_type.clear()
        i = sum = 0
        for key in temp.keys():
            if i < 8:
                column_type[key] = temp[key]
            else:
                sum += temp[key]
            i += 1
        column_type['others'] = sum
    labels = list(column_type.keys())
    sizes = list(column_type.values())
    plt.figure(figsize=(20, 15), dpi=80)
    plt.pie(sizes,
            labels=labels,
            autopct='%3.2f%%',
            shadow=False,
            startangle=180
            )
    plt.legend()
    plt.savefig(r'./' + dir + r'/' +
                column_name.replace('/', 'Or') + '.png')


# 世界地图


def draw_map(data, dir):
    plt.figure(figsize=(40, 30))
    m = Basemap()
    m.drawcoastlines()
    m.drawcountries(color='grey', linewidth=2)
    m.drawstates()
    m.drawmapboundary(fill_color='#87CEEB')
    m.fillcontinents(color='green', lake_color='#87CEEB')
    # 经纬线
    m.drawparallels(circles=np.linspace(-90, 90, 7),
                    labels=[1, 0, 0, 0], color='gray')
    m.drawmeridians(meridians=np.linspace(-180, 180, 13),
                    labels=[0, 0, 0, 1], color='pink')
    # 两数据集经纬叫法不一致
    if dir == 'starbuck':
        lon = data['Longitude']
        lat = data['Latitude']
    else:
        lon = data['lng']
        lat = data['lat']
    # 打点
    m.scatter(lon, lat, s=2, c='red')
    plt.savefig(r'./' + dir + r'/' + 'Worldmap' + '.png')
